`include "defines.v"
module rx_fsm(
    input [1:0] rx_rate_i,
    input dewhiten_bypass_i,
    input [7:0] rx_up_us_i,
    output reg [7:0] rx_byte_o,
    output [23:0] rx_crc_rslt_o,
    output reg [23:0] rx_crc_raw_o,
    output reg rx_mic_err_o,
    input aes_ccm_en_i,
    output reg rx_fifo_push_o,
    input start_i,
    output reg rx_en_o,
    input rx_data_i,
    input rx_data_valid_i,
    input sync_found_i,

    output reg [7:0] rx_payload_length_o,
    output reg aes_ccm_start_o,
    output reg [7:0] encrypted_byte_o,
    output [7:0] encrypted_byte_cnt_o,
    output reg encrypted_byte_update_o,
    output reg [7:0] ccm_aad_o,
    input decrypted_byte_valid_i,
    input [31:0] rx_mic_calc_i,
    input [7:0] decrypted_byte_i,

    output crc_data_o,
    output rx_crc_en_o,
    input [23:0] rx_crc_calc_i,

    output whitened_data_o,
    output rx_whiten_shift_o,
    input dewhitened_data_i,

    input clk,
    input rst_n
);

reg [2:0] c_st;
reg [2:0] n_st;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) c_st<= `RX_IDLE;
    else c_st <= n_st;
end

reg [4:0] cnt;
wire sym_end;
assign sym_end = (rx_rate_i == 0 && cnt == `CLK_MHZ - 1) || (rx_rate_i == 1 && cnt == `CLK_MHZ/2 - 1);
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) begin
        cnt <= 0;
        rx_en_o <= 0;
    end
    else if(n_st != `RX_IDLE) begin
        rx_en_o <= 1;
        if(sym_end) cnt <= 0;
        else cnt <= cnt + 1;
    end
    else begin
        rx_en_o <= 0;
        cnt <= 0;
    end
end

reg [8:0] up_cnt;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) up_cnt <= 0;
    else if(c_st == `RX_UP) begin
        if(sym_end) up_cnt <= up_cnt + 1;
    end
    else up_cnt <= 0;
end

reg [3:0] header_cnt;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) header_cnt <= 0;
    else if(c_st == `RX_HEADER) begin
        if(rx_data_valid_i) header_cnt <= header_cnt + 1;
    end
    else header_cnt <= 0;
end

reg [7:0] dewhitened_pdu_byte;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) begin 
        rx_payload_length_o <= 0;
        aes_ccm_start_o <= 0;
    end
    else if(c_st == `RX_JUDGE) begin 
        rx_payload_length_o <= dewhitened_pdu_byte;
        aes_ccm_start_o <= aes_ccm_en_i && dewhitened_pdu_byte;
    end
    else if(aes_ccm_start_o) begin
        aes_ccm_start_o <= 0;
    end
end

always @(posedge clk,negedge rst_n) begin
    if(!rst_n) ccm_aad_o<= 0;
    else if(header_cnt == 8) ccm_aad_o <= dewhitened_pdu_byte;
end

reg [10:0] payload_cnt;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) payload_cnt <= 0;
    else if(c_st == `RX_PAYLOAD) begin
        if(rx_data_valid_i) payload_cnt <= payload_cnt + 1;
    end
    else payload_cnt <= 0;
end

reg [4:0] crc_cnt;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) crc_cnt <= 0;
    else if(c_st == `RX_CRC) begin
        if(rx_data_valid_i) crc_cnt <= crc_cnt + 1;
    end
    else crc_cnt <= 0;
end

assign whitened_data_o = rx_data_i;
assign rx_whiten_shift_o = rx_data_valid_i;
assign crc_data_o = dewhiten_bypass_i ? rx_data_i : dewhitened_data_i;
assign rx_crc_en_o = rx_data_valid_i;
assign rx_crc_rslt_o = rx_crc_calc_i;
reg [31:0] mic_rx;
wire [5:0] mic_cnt;
wire on_mic_rx;
assign mic_cnt = 8*rx_payload_length_o - payload_cnt - 1;
assign on_mic_rx = 8*rx_payload_length_o - payload_cnt - 1 <= 31 && aes_ccm_en_i && rx_payload_length_o;

reg dewhitened_byte_ready;
reg [8:0] dewhitened_byte_cnt;
wire encrypted_byte_receiving;
assign encrypted_byte_cnt_o = dewhitened_byte_cnt - 3;
assign encrypted_byte_receiving = dewhitened_byte_cnt>2&&dewhitened_byte_cnt-2<=rx_payload_length_o-4;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) begin 
        encrypted_byte_o <= 0;
        encrypted_byte_update_o <= 0;
    end
    else if(encrypted_byte_update_o) encrypted_byte_update_o <= 0;
    else if(dewhitened_byte_ready&&encrypted_byte_receiving) begin
        encrypted_byte_o <= dewhitened_pdu_byte;
        encrypted_byte_update_o <= 1;
    end
end

always @(posedge clk,negedge rst_n) begin
    if(!rst_n) begin
        rx_fifo_push_o <= 0;
        rx_byte_o <= 0;
    end
    else begin
        if(rx_fifo_push_o) rx_fifo_push_o <= 0;
        if(aes_ccm_en_i && encrypted_byte_receiving) begin
            if(decrypted_byte_valid_i) begin
                rx_fifo_push_o <= 1;
                rx_byte_o <= decrypted_byte_i;
            end
        end
        else begin
            if(dewhitened_byte_ready) begin
                rx_fifo_push_o <= 1;
                rx_byte_o <= dewhitened_pdu_byte;
            end
        end
    end
end
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) begin
        dewhitened_pdu_byte <= 0;
        mic_rx <= 0;
        rx_crc_raw_o <= 0;
        dewhitened_byte_ready <= 0;
        dewhitened_byte_cnt <= 0;
    end
    else begin
        if(dewhitened_byte_ready) dewhitened_byte_ready <= 0;
        case (c_st)
            `RX_HEADER: begin
                if(rx_data_valid_i) begin
                    dewhitened_pdu_byte[header_cnt%8] <= crc_data_o;
                    if(header_cnt%8 == 7) begin
                        dewhitened_byte_ready <= 1;
                        dewhitened_byte_cnt <= dewhitened_byte_cnt + 1;
                    end
                end
            end
            `RX_JUDGE: begin

            end
            `RX_PAYLOAD: begin
                if(rx_data_valid_i) begin
                    dewhitened_pdu_byte[payload_cnt%8] <= crc_data_o;
                    if(payload_cnt%8 == 7) begin
                        dewhitened_byte_ready <= 1;
                        dewhitened_byte_cnt <= dewhitened_byte_cnt + 1;
                    end
                    if(aes_ccm_en_i && on_mic_rx) mic_rx[31 - mic_cnt] = crc_data_o;
                end
            end
            `RX_CRC: if(rx_data_valid_i) begin
                rx_crc_raw_o[23 - crc_cnt] <= crc_data_o;
            end
            default: begin
                dewhitened_pdu_byte <= 0;
                dewhitened_byte_cnt <= 0;
            end
        endcase
    end
end

always @(posedge clk,negedge rst_n) begin
    if(!rst_n) rx_mic_err_o <= 0;
    else if(c_st == `RX_CRC) begin
        if(aes_ccm_en_i && rx_payload_length_o && mic_rx != rx_mic_calc_i) rx_mic_err_o <= 1;
        else rx_mic_err_o <= 0;
    end
end

always @(*) begin
    n_st = c_st;
    case (c_st)
        `RX_IDLE:    if(start_i) n_st = `RX_UP;
        `RX_UP:      if((rx_rate_i ? (up_cnt == 2*rx_up_us_i + 1): (up_cnt == rx_up_us_i)) && sym_end) n_st = `RX_WINDOW;
        `RX_WINDOW:  if(sync_found_i) n_st = `RX_HEADER;
        `RX_HEADER:  if(rx_data_valid_i && header_cnt == 15) n_st = `RX_JUDGE;
        `RX_JUDGE:   if(dewhitened_pdu_byte) n_st = `RX_PAYLOAD;
                    else n_st = `RX_CRC;
        `RX_PAYLOAD: if(rx_data_valid_i && payload_cnt + 1 == 8*rx_payload_length_o) n_st = `RX_CRC;
        `RX_CRC:     if(rx_data_valid_i && crc_cnt == 23) n_st = `RX_IDLE;
    endcase
end
endmodule